import random
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from albumentations.pytorch.transforms import ToTensorV2
import matplotlib.pyplot as plt
import h5py
from torchvision.utils import make_grid
from scipy.ndimage.interpolation import zoom
import itertools
from torch.utils.data.sampler import Sampler
from torchvision import transforms
import torch
from scipy import ndimage


class ACDC_Semi(Dataset):

    PALETTE = np.array([
        [0, 0, 0],
        [0, 0, 255],
        [0, 255, 0],
        [255, 0, 0],
    ])

    def __init__(self, root=r"E:\note\ssl\data\ACDC", split="train", transform=None):

        super(ACDC_Semi, self).__init__()
        self.split = split
        self.root = root
        self.transform = transform
        self.sample_list = []
        self.load_annotations()  # 加载文件路径
        print("total {} samples".format(len(self.sample_list)))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, idx):
        case = self.sample_list[idx]
        h5f = h5py.File(case, "r")

        image = np.array(h5f["image"][:],dtype=np.float32)
        mask = np.array(h5f["label"][:],dtype=np.uint8)

        if self.transform is not None and self.split == "train":
            result = self.transform(image=image, mask=mask)
            image = result["image"]
            mask = result["mask"]

        return image, mask

    def label_to_img(self, label):
        if isinstance(label, torch.Tensor):
            label = label.numpy()
        if not isinstance(label, np.ndarray):
            label = np.array(label)
        label = label.astype(np.uint8)
        label[label == 255] = 0
        img = self.PALETTE[label]
        if len(img.shape) == 4:
            img = torch.tensor(img).permute(0, 3, 1, 2)
            img = make_grid(tensor=img, nrow=2, scale_each=True)
            img = img.permute(1, 2, 0).numpy()

        return img.astype(np.uint8)

    def load_annotations(self):
        if self.split == "train":
            with open(self.root + "/train_slices.list", "r") as f1:
                self.sample_list = f1.readlines()
            self.sample_list = [item.replace("\n", "") for item in self.sample_list]
            self.sample_list = [self.root + "/data/slices/{}.h5".format(item) for item in self.sample_list]
        elif self.split == "val":
            with open(self.root + "/val.list", "r") as f:
                self.sample_list = f.readlines()
            self.sample_list = [item.replace("\n", "") for item in self.sample_list]
            self.sample_list = [self.root + "/data/{}.h5".format(item) for item in self.sample_list]
        else:
            with open(self.root + "/test.list", "r") as f:
                self.sample_list = f.readlines()
            self.sample_list = [item.replace("\n", "") for item in self.sample_list]
            self.sample_list = [self.root + "/data/{}.h5".format(item) for item in self.sample_list]
            
        self.sample_list=np.array(self.sample_list)


def random_rot_flip(image, label):
    k = np.random.randint(0, 4)
    image = np.rot90(image, k)
    label = np.rot90(label, k)
    axis = np.random.randint(0, 2)
    image = np.flip(image, axis=axis).copy()
    label = np.flip(label, axis=axis).copy()
    return image, label


def random_rotate(image, label):
    angle = np.random.randint(-20, 20)
    image = ndimage.rotate(image, angle, order=0, reshape=False)
    label = ndimage.rotate(label, angle, order=0, reshape=False)
    return image, label


class RandomGenerator(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, image,mask):
        if random.random() > 0.5:
            image, mask = random_rot_flip(image, mask)
        elif random.random() > 0.5:
            image, mask = random_rotate(image, mask)
        x, y = image.shape
        image = zoom(
            image, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        mask = zoom(
            mask, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        image = torch.from_numpy(
            image.astype(np.float32)).unsqueeze(0)
        mask = torch.from_numpy(mask.astype(np.uint8))
        sample = {'image': image, 'mask': mask}
        return sample


class TwoStreamBatchSampler(Sampler):
    """Iterate two sets of indices

    An 'epoch' is one iteration through the primary indices.
    During the epoch, the secondary indices are iterated through
    as many times as needed.
    """

    def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size):
        self.primary_indices = primary_indices
        self.secondary_indices = secondary_indices
        self.secondary_batch_size = secondary_batch_size
        self.primary_batch_size = batch_size - secondary_batch_size

        assert len(self.primary_indices) >= self.primary_batch_size > 0
        assert len(self.secondary_indices) >= self.secondary_batch_size > 0

    def __iter__(self):
        primary_iter = iterate_once(self.primary_indices)
        secondary_iter = iterate_eternally(self.secondary_indices)
        return (
            primary_batch + secondary_batch
            for (primary_batch, secondary_batch)
            in zip(grouper(primary_iter, self.primary_batch_size),
                   grouper(secondary_iter, self.secondary_batch_size))
        )

    def __len__(self):
        return len(self.primary_indices) // self.primary_batch_size


def iterate_once(iterable):
    return np.random.permutation(iterable)


def iterate_eternally(indices):
    def infinite_shuffles():
        while True:
            yield np.random.permutation(indices)
    return itertools.chain.from_iterable(infinite_shuffles())


def grouper(iterable, n):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3) --> ABC DEF"
    args = [iter(iterable)] * n
    return zip(*args)

def patients_to_slices(patiens_num):
    if patiens_num == 0.05:
        return 68
    elif patiens_num == 0.1:
        return 136
    elif patiens_num == 0.2:
        return 256
    else:
        ref_dict = {"3": 68, "7": 136,
                        "14": 256, "21": 396, "28": 512, "35": 664, "140": 824, "300": 1024}
        
        return ref_dict[str(patiens_num)]

def worker_init_fn(worker_id):
        random.seed(1337 + worker_id)

def get_ssl_acdc_semi_loader(root=r'/home/ubuntu/data/ACDC', batch_size=8, unlabel_batch_size=24, train_crop_size=(224, 224), label_num=0.2):
    """
    :param root: 数据集路径
    :param batch_size: 有标注数据批次大小
    :param unlabel_batch_size: 无标注数据的batch大小
    :param label_num: 有标签的数量
    :return:
    """
    train_transform=RandomGenerator(train_crop_size)
    train_dataset = ACDC_Semi(root=root, split="train", transform=train_transform)
    test_dataset = ACDC_Semi(root=root, split="test")

    total_slices = len(train_dataset)

    labeled_slice = patients_to_slices(label_num)

    labeled_idxs = list(range(0, labeled_slice))
    unlabeled_idxs = list(range(labeled_slice, total_slices))
    batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size + unlabel_batch_size, unlabel_batch_size)

    trainloader = DataLoader(train_dataset, batch_sampler=batch_sampler,num_workers=4, worker_init_fn=worker_init_fn)
    test_loader = DataLoader(test_dataset, batch_size=1, num_workers=4, shuffle=False)
    return trainloader, test_loader

def show(im):
    im = im.numpy().squeeze()
    fig=plt.figure()
    plt.imshow(im, cmap="gray")
    plt.show()
    fig.savefig("result.png")


def show_label(mask, path="label.jpg"):
    plt.figure()
    plt.imshow(mask)
    plt.show()
    Image.fromarray(mask).save(path)

# def get_ssl_acdc_loader(root=r'/home/ubuntu/data/ACDC', batch_size=8, unlabel_batch_size=24, train_crop_size=(224, 224), label_num=0.2,seed=1337):
#     """
#     :param root: 数据集路径
#     :param batch_size: 有标注数据批次大小
#     :param unlabel_batch_size: 无标注数据的batch大小
#     :param label_num: 有标签的数量
#     :return:
#     """

#     def worker_init_fn(worker_id):
#         random.seed(seed + worker_id)
    
#     train_transform=RandomGenerator(train_crop_size)
#     train_dataset = ACDC(root=root, split="train", transform=train_transform)
#     total_slices=len(train_dataset)
#     labeled_slice=patients_to_slices("ACDC",label_num)
#     labeled_idxs = list(range(0, labeled_slice))
#     unlabeled_idxs = list(range(labeled_slice, total_slices))

#     batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size+unlabel_batch_size,unlabel_batch_size)
#     train_loader=DataLoader(train_dataset,batch_sampler=batch_sampler,num_workers=4,pin_memory=True,worker_init_fn=worker_init_fn)

#     test_dataset = ACDC(root=root, split="test")
#     test_loader = DataLoader(test_dataset, batch_size=1, num_workers=4, shuffle=False)
#     return train_loader, test_loader



if __name__ == '__main__':

    train_dataloader, test_dataloader = get_acdc_loader()
    # print(len(train_dataloader))
    # print(len(test_dataloader))
    # print(len(test_dataloader.dataset))
    for image, label in train_dataloader:
        print(image.shape)
        print(label.shape)
        print(np.unique(label.numpy()))
        show(image[0])
        show_label(train_dataloader.dataset.label_to_img(label))
        break

    for sample in test_dataloader:
        image, label = sample
        print(image.shape)
        print(label.shape)
        print(np.unique(label.numpy()))
        # show(image[0])
        # show_label(label[0].numpy())
        break
